import numpy as np
import matplotlib.pyplot as plt
import math

plt.style.use('default')
plt.rcParams['font.family'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 1. 定义混合波形
srate = 3000
t = np.linspace(0, 1, srate)
s1 = 1.5 * np.sin(-2 * np.pi * 30 * t)
s2 = 3 * np.cos(-2 * np.pi * 5 * t)
s3 = 10 * np.sin(-2 * np.pi * 100 * t)
s4 = 20 * np.cos(-2 * np.pi * 120 * t)
m = s1 + s2 + s3 + s4
n = m + 2.5 * np.random.randn(len(t))

fcoefs = np.fft.fft(m)
amp_spec = 2 * np.abs(fcoefs / srate)

freqs = np.fft.fftfreq(len(m), 1 / srate) 
amp_spec_smoothed = np.copy(amp_spec) 
amp_spec_smoothed[amp_spec < 0.05 * np.max(amp_spec)] = 0 
amp_spec_smoothed[(np.abs(freqs) > 110)] = 0 

freq_shifted = np.fft.fftshift(freqs)
amp_shifted = np.fft.fftshift(amp_spec) 
amp_shifted_smoothed = np.fft.fftshift(amp_spec_smoothed) 


plt.subplot(1, 2, 1)
plt.stem(freq_shifted, amp_shifted)
plt.xlim(-150, 150)
plt.xlabel('频率')
plt.ylabel('幅值')
plt.grid()

plt.subplot(1, 2, 2)
plt.stem(freq_shifted, amp_shifted_smoothed)
plt.xlim(-150, 150)
plt.ylim(-1,21)
plt.xlabel('频率')
plt.ylabel('幅值')
plt.grid()

plt.tight_layout

plt.show()